/**
 * 
 */
package edu.arizona.bio5.onto.decomposition;

import java.net.URI;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.semanticweb.owlapi.model.IRI;
import org.semanticweb.owlapi.model.OWLAxiom;
import org.semanticweb.owlapi.model.OWLEntity;

import uk.ac.manchester.cs.graph.DirectedGraph;
import uk.ac.manchester.cs.graph.DirectedGraphImpl;
import uk.ac.manchester.cs.graph.Edge;
import uk.ac.manchester.cs.graph.alg.BFS;
import uk.ac.manchester.cs.graph.alg.ClosureVisitor;
import uk.ac.manchester.cs.graph.alg.VertexVisitor;
import uk.ac.manchester.cs.owlapi.modularity.ModuleType;

/**
 * Stores an atomic decomposition of an ontology
 * 
 * @author Pavel Klinov
 *
 * pklinov@cs.man.ac.uk, pklinov@clarkparsia.com
 * 
 * Mar 5, 2011
 */
public class AtomicDecomposition {
	
	private int m_idCnt = 1;
	
	protected DirectedGraph<Atom, Object> m_adGraph = new DirectedGraphImpl<Atom, Object>();
	//Inverted graph for bottom-up traversal
	//FIXME Get rid of this since the directed graph now stores both outgoing and incoming edges
	protected DirectedGraph<Atom, Object> m_adInvGraph = new DirectedGraphImpl<Atom, Object>();
	protected Map<OWLAxiom, Atom> m_axiomBasedIndex = new HashMap<OWLAxiom, Atom>();
	/*
	 * Maps each term to a set of atoms s.t. each atom could be potentially relevant to the term.
	 * This index is built externally as it depends on the notion of labels
	 */
	protected Map<URI, Set<Atom>> m_termBasedIndex = null;
	protected Set<OWLAxiom> m_globalAxioms = null;
	protected Set<OWLAxiom> m_tautologies = null;
	protected ModuleType m_modType = null;
	private URI m_ontoURI = null;
	
	
	public AtomicDecomposition(URI ontoURI, ModuleType mType) {
		m_ontoURI = ontoURI;
		m_modType = mType;
	}
	
	public URI getOntologyURI() {
		return m_ontoURI;
	}

	public void setGlobalAxioms(Set<OWLAxiom> axioms) {
		m_globalAxioms = axioms;
	}
	
	public Set<OWLAxiom> getGlobalAxioms() {
		return m_globalAxioms;
	}
	
	public void setTautologies(Set<OWLAxiom> axioms) {
		m_tautologies = axioms;
	}
	
	public ModuleType getModuleType() {
		
		return m_modType;
	}
	
	protected void setModuleType(ModuleType mType) {
		
		m_modType = mType;
	}
	
	public DirectedGraph<Atom, Object> getGraph() {
		
		return m_adGraph;
	}
	
	public Set<Atom> getAtoms() {
		
		return m_adGraph.getVertices();
	}
	
	public Atom getAtomForAxiom(OWLAxiom axiom) {
		
		return m_axiomBasedIndex.get( axiom );
	}
	
	public Collection<Atom> getAtomsForTerm(URI termURI) {
		//A stupid implementation, use only if can't use the term index/labels for whatever reason
		Collection<Atom> atoms = new ArrayList<Atom>();
		IRI termIRI = IRI.create(termURI);
		
		for (Atom atom : getAtoms()) {
			for (OWLEntity entity : atom.getSignature()) {
				if (entity.getIRI().equals(termIRI)) {
					atoms.add(atom);
				}
			}
		}
		
		return atoms;
	}
	
	public Map<URI, Set<Atom>> getTermBasedIndex() {
		
		return m_termBasedIndex;
	}
	
	public void setTermBasedIndex(Map<URI, Set<Atom>> tbIndex) {
		
		m_termBasedIndex = tbIndex;
	}
	
	public MemAtom addAtom(Set<OWLAxiom> axioms) {
		
		MemAtom atom = new MemAtom(m_idCnt++, axioms);
		
		m_adGraph.addVertex( atom );
		m_adInvGraph.addVertex( atom );
		
		updateAxiomIndex(atom);
		
		return atom;
	}
	
	private void updateAxiomIndex(Atom atom) {
		
		for (OWLAxiom axiom : atom.getAxioms()) m_axiomBasedIndex.put(axiom, atom);
	}

	/*
	 * Returns top atoms that have this atom in their principal ideals
	 */
	public Set<Atom> getTopAtoms(Atom atom) {
		
		TopAtomVisitor visitor = new TopAtomVisitor();
		
		new BFS().search(m_adInvGraph, atom, visitor);
		
		return visitor.getAtoms();
	}

	public boolean isTopAtom(Atom atom) {
		
		return m_adInvGraph.getAdjacentVertices(atom).isEmpty();
	}
	
	public boolean isBottomAtom(Atom atom) {
		
		return m_adGraph.getAdjacentVertices(atom).isEmpty();
	}

	public void addAxiomToAtom(OWLAxiom axiom, MemAtom atom) {
		
		if (m_adGraph.getVertices().contains(atom)) {
			
			atom.addAxiom(axiom);
			m_axiomBasedIndex.put(axiom, atom);
			
		} else throw new IllegalArgumentException("Atom not found");
	}
	
	public void addDependency(Atom dependentAtom, Atom dependentUponAtom) {
		if (m_adGraph.findEdge(dependentAtom, dependentUponAtom) == null) { 
			m_adGraph.addEdge( "", dependentAtom, dependentUponAtom );
			m_adInvGraph.addEdge( "", dependentUponAtom, dependentAtom );
		}
	}
	
	protected void removeAtom(Atom atom) {
		
		m_adGraph.removeVertex(atom);
		m_adInvGraph.removeVertex(atom);
		
		for (OWLAxiom axiom : atom.getAxioms()) {
			m_axiomBasedIndex.remove(axiom);
		}
	}
	
	/*
	 * Merges beta into alpha
	 */
	public void mergeAtoms(MemAtom alpha, Atom beta) {
		//update dependencies
		for (Atom dep : getDependents(beta, true)) {
			addDependency(dep, alpha);
		}
		
		for (Atom dep : getDependencies(beta, true)) {
			addDependency(alpha, dep);
		}
		
		for (OWLAxiom axiom : beta.getAllAxioms()) {
			addAxiomToAtom(axiom, alpha);
		}
		
		removeAtom(beta);
	}
	
	protected void removeDependency(Atom dependentAtom, Atom dependentUponAtom) {
		
		Edge<Atom, Object> direct = m_adGraph.findEdge(dependentAtom, dependentUponAtom );
		Edge<Atom, Object> reversed = m_adInvGraph.findEdge(dependentUponAtom, dependentAtom );
		
		if (direct != null) m_adGraph.removeEdge(direct);
		if (reversed != null) m_adInvGraph.removeEdge(reversed);
	}

	/*
	 * Returns all axioms in this atom and all atoms small than this
	 */
	public AxiomSet getPrincipalIdeal(Atom atom) {
		
		Set<OWLAxiom> axioms = Collections.emptySet();

		if (m_adGraph.getVertices().contains(atom)) {
			// Run breadth-first search here
			AtomVisitor visitor = new AtomVisitor();

			new BFS().search(m_adGraph, atom, visitor);
			axioms = visitor.getAxioms();
		}

		return new AxiomSet(axioms);
	}
	
	/*
	 * TODO Optimize
	 */
	public Set<OWLEntity> getPrincipalIdealSignature(Atom atom) {
		
		if (isBottomAtom(atom)) {
			
			return atom.getSignature();
			
		} else {
			
			Set<OWLEntity> piSig = new HashSet<OWLEntity>(atom.getSignature());
			
			for (Atom lower : getDependencies(atom, true)) piSig.addAll(getPrincipalIdealSignature(lower));
			
			return piSig;
		}
	}
	
	/*
	 * Returns the set of atoms the given atom depends on. The set contains the given atom.
	 */
	public Set<Atom> getDependencies(Atom atom) {
		
		ClosureVisitor<Atom> visitor = new ClosureVisitor<Atom>();
		
		new BFS().search(m_adGraph, atom, visitor);
		
		return visitor.getVertices();
	}
	
	public Set<Atom> getDependencies(Atom atom, boolean direct) {
		
		return direct ? m_adGraph.getAdjacentVertices(atom) : getDependencies(atom);
	}	
	
	
	/*
	 * Returns the set of atoms that depend on the given atom depends on. The set contains the given atom.
	 */
	public Set<Atom> getDependents(Atom atom) {
		
		ClosureVisitor<Atom> visitor = new ClosureVisitor<Atom>();
		
		new BFS().search(m_adInvGraph, atom, visitor);
		
		return visitor.getVertices();
	}
	
	public Set<Atom> getDependents(Atom atom, boolean direct) {
		
		return direct ? m_adInvGraph.getAdjacentVertices(atom) : getDependents(atom);
	}	
	
	/*
	 * Returns the connected component for the given atom
	 */
	public Set<Atom> getRelatedAtoms(Atom atom) {
		
		Set<Atom> cc = new HashSet<Atom>();
		Set<Atom> reachedAtoms = new HashSet<Atom>();
		Set<Atom> newAtoms = new HashSet<Atom>();
		
		reachedAtoms.add(atom);
		
		while(!reachedAtoms.isEmpty()) {
			
			for (Atom reached : reachedAtoms) {
				
				for (Atom lower : getDependencies(reached)) {
					
					if (cc.add(lower)) newAtoms.add(lower);
				}
				
				for (Atom higher : getDependents(reached)) {
					
					if (cc.add(higher)) newAtoms.add(higher);
				}
			}
			
			reachedAtoms.clear();
			reachedAtoms.addAll(newAtoms);
			newAtoms.clear();
		}
		
		return cc;
	}
	
	public ADStats stats() {
		//TODO Cache the stats for future reference
		ADStats stats = new ADStats();
		//Compute the stats
		int largest = 0;
		int totalSize = 0;
		int num = m_adGraph.getVertices().size();
		int maxPISizeAtoms = 0;
		int totalPISizeAtoms = 0;
		int maxPISizeAxioms = 0;
		int totalPISizeAxioms = 0;		
		int ccNum = 0;
		Atom biggestAtom = null;
		Set<Atom> ccAtoms = new HashSet<Atom>(m_adGraph.getVertices().size());
		
		for (Atom atom : m_adGraph.getVertices()) {
			
			AxiomSet pi = getPrincipalIdeal(atom);
			Set<Atom> deps = getDependencies(atom);
			
			if (atom.getAxioms().size() > largest) {
				
				largest = atom.getAxioms().size();
				biggestAtom = atom;
			}
			
			maxPISizeAxioms = Math.max(maxPISizeAxioms, pi.getAxioms().size());
			maxPISizeAtoms = Math.max(maxPISizeAtoms, deps.size());
			
			totalSize += atom.getAxioms().size();
			totalPISizeAxioms += pi.getAxioms().size();
			totalPISizeAtoms += deps.size();

			if (!ccAtoms.contains(atom)){
				ccAtoms.addAll(getRelatedAtoms(atom));
				ccNum++;
			}
		}
		
		stats.setBiggestAtom(biggestAtom);
		stats.setAtoms(m_adGraph.getVertices().size());
		stats.setBiggestAtomSize(largest);
		stats.setAvgAtomSize((1d * totalSize / num));
		stats.setMaxPISizeInAtoms(maxPISizeAtoms);
		stats.setMaxPISizeInAxioms(maxPISizeAxioms);
		stats.setAvgPISizeInAtoms(1d*totalPISizeAtoms / num);
		stats.setAvgPISizeInAxioms(1d*totalPISizeAxioms / num);
		stats.setConnectedComponentsNumber(ccNum);
		stats.setTotalNumberOfAxioms(totalSize);
		
		return stats;
	}
	
	/*
	 * Basic term index: each term is mapped to the set of atoms where it appears
	 */
	public void buildBasicTermIndex() {
		Map<URI, Set<Atom>> tbIndex = new HashMap<URI, Set<Atom>>(getAtoms().size() * 5);
		
		for (Atom atom : getAtoms()) {
			for (OWLEntity term : atom.getSignature()) {
				Set<Atom> atomsForTerm = tbIndex.get(term.getIRI().toURI());
				
				atomsForTerm = atomsForTerm == null ? new HashSet<Atom>() : atomsForTerm;
				atomsForTerm.add(atom);
				tbIndex.put(term.getIRI().toURI(), atomsForTerm);
			}
		}
		
		setTermBasedIndex(tbIndex);
	}

	public Set<Atom> getTopAtoms() {
		
		return m_adGraph.getRoots();
	}
	
	public Set<Atom> getBottomAtoms() {
		
		return m_adInvGraph.getRoots();
	}
	
	public String toString() {
		//Currently prints only atoms
		//TODO Print the graph
		return m_adGraph.getVertices().toString();
	}

	public Atom getAtomByID(Object id) {
		// FIXME Inefficient sequential search
		for (Atom atom : getAtoms()) {
			
			if (atom.getID().toString().equals(id.toString())) return atom;
		}
		
		return null;
	}
	
	class AtomVisitor implements VertexVisitor<Atom> {

		private Set<OWLAxiom> m_closure = new HashSet<OWLAxiom>();
		
		@Override
		public boolean visit(Atom vertex) {

			m_closure.addAll( vertex.getAxioms() );
			
			return true;
		}
		
		public Set<OWLAxiom> getAxioms() {
			
			return m_closure;
		}
	}
	
	class TopAtomVisitor implements VertexVisitor<Atom> {

		private Set<Atom> m_topAtoms = new HashSet<Atom>();
		
		@Override
		public boolean visit(Atom vertex) {

			if (isTopAtom(vertex)) m_topAtoms.add(vertex);
			
			return true;
		}
		
		public Set<Atom> getAtoms() {
			
			return m_topAtoms;
		}
	}
}